package com.tyylab.kestrel.framework.util;

import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;

/**
 * Context utility for managing thread-local variables
 * Supports tenant_id, request_id and other header propagation
 * 
 * @author tyylab
 */
@Slf4j
public class ContextUtil {
    
    private static final ThreadLocal<Map<String, Object>> CONTEXT = new ThreadLocal<>();
    
    // Default header keys
    public static final String TENANT_ID = "tenant_id";
    public static final String REQUEST_ID = "request_id";
    
    /**
     * Set a value to thread local context
     * @param key the key
     * @param value the value
     */
    public static void set(String key, Object value) {
        Map<String, Object> context = getContext();
        context.put(key, value);
        log.debug("Set context key: " + key + ", value: " + value);
    }
    
    /**
     * Get a value from thread local context
     * @param key the key
     * @return the value, null if not found
     */
    public static Object get(String key) {
        Map<String, Object> context = getContext();
        Object value = context.get(key);
        log.debug("Get context key: " + key + ", value: " + value);
        return value;
    }
    
    /**
     * Get a value from thread local context with default value
     * @param key the key
     * @param defaultValue the default value
     * @return the value or default value if not found
     */
    public static Object get(String key, Object defaultValue) {
        Object value = get(key);
        return value != null ? value : defaultValue;
    }
    
    /**
     * Get tenant ID from context
     * @return tenant ID, null if not found
     */
    public static String getTenantId() {
        Object value = get(TENANT_ID);
        return value != null ? value.toString() : null;
    }
    
    /**
     * Set tenant ID to context
     * @param tenantId the tenant ID
     */
    public static void setTenantId(String tenantId) {
        set(TENANT_ID, tenantId);
    }
    
    /**
     * Get request ID from context
     * @return request ID, null if not found
     */
    public static String getRequestId() {
        Object value = get(REQUEST_ID);
        return value != null ? value.toString() : null;
    }
    
    /**
     * Set request ID to context
     * @param requestId the request ID
     */
    public static void setRequestId(String requestId) {
        set(REQUEST_ID, requestId);
    }
    
    /**
     * Remove a key from context
     * @param key the key to remove
     */
    public static void remove(String key) {
        Map<String, Object> context = getContext();
        context.remove(key);
        log.debug("Remove context key: " + key);
    }
    
    /**
     * Clear all context
     */
    public static void clear() {
        CONTEXT.remove();
        log.debug("Clear all context");
    }
    
    /**
     * Get all context as a copy
     * @return a copy of all context
     */
    public static Map<String, Object> getAll() {
        Map<String, Object> context = getContext();
        return new HashMap<>(context);
    }
    
    /**
     * Get the internal context map, create if not exists
     * @return the context map
     */
    private static Map<String, Object> getContext() {
        Map<String, Object> context = CONTEXT.get();
        if (context == null) {
            context = new HashMap<>();
            CONTEXT.set(context);
        }
        return context;
    }
} 